'''
Author: zhanwei xu
Date: 2023-05-27 17:12:07
LastEditors: zhanwei xu
LastEditTime: 2023-10-27 15:37:55
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''
from fastapi import APIRouter, Request, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import random
import string
import os
from io import BytesIO
from .visionbot import Visionbot
from .geminibot import Chatbot
vision_router = APIRouter()
visionbot = Visionbot()
geminibot = Chatbot()
def generate_random_string(length: int = 16) -> str:
    letters = string.ascii_letters + string.digits
    return ''.join(random.choice(letters) for i in range(length))

def resize_image(image: Image, max_size: int = 1200) -> Image:
    width, height = image.size
    if width > max_size or height > max_size:
        if width > height:
            new_height = int(max_size * (height / width))
            new_width = max_size
        else:
            new_width = int(max_size * (width / height))
            new_height = max_size
        return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    return image

@vision_router.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
    # Ensure the file is an image
    if not file.content_type.startswith('image/'):
        return JSONResponse(status_code=400, content={"message": "File is not an image"})

    # Generate random file name with a ".jpg" extension
    random_name = generate_random_string()
    file_dir = '/var/www/html/images/'
    file_name = f'{file_dir}{random_name}.jpg'

    # Read image file
    image_data = await file.read()
    image = Image.open(BytesIO(image_data))
    
    # Resize the image
    resized_image = resize_image(image)
    
    # 如果图片是带有透明通道的 RGBA，转换为 RGB 并使用白色填充透明
    if resized_image.mode == 'RGBA':
        background = Image.new('RGB', resized_image.size, (255, 255, 255))
        background.paste(resized_image, mask=resized_image.split()[3])  # 3 是 alpha 通道的索引
        resized_image = background
    # Check if the image is in palette mode ('P') and convert it to RGB mode
    elif resized_image.mode == 'P':
        resized_image = resized_image.convert('RGB')
    # Save the resized image in JPEG format regardless of the original format
    resized_image.save(file_name, "JPEG")

    return {"filename": random_name}

@vision_router.post("/vision_chat",summary="return the response of the message using vision")
async def vision_chat(request: Request):
    return await visionbot.visionchat(request)
@vision_router.post("/vision_gemini",summary="return the response of the message using vision")
async def vision_chat(request: Request):
    return await geminibot.chatbot(request)
